This tutorial illustrates the core visualization utilities available in Ax.
import numpy as np
from ax.service.ax_client import AxClient
from ax.modelbridge.cross_validation import cross_validate
from ax.plot.contour import interact_contour
from ax.plot.diagnostic import interact_cross_validation
from ax.plot.scatter import(
interact_fitted,
plot_objective_vs_constraints,
tile_fitted,
)
from ax.plot.slice import plot_slice
from ax.utils.measurement.synthetic_functions import hartmann6
from ax.utils.notebook.plotting import render, init_notebook_plotting
init_notebook_plotting()
[INFO 02-02 05:17:40] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.
The vizualizations require an experiment object and a model fit on the evaluated data. The routine below is a copy of the Service API tutorial, so the explanation here is omitted. Retrieving the experiment and model objects for each API paradigm is shown in the respective tutorials
noise_sd = 0.1
param_names = [f"x{i+1}" for i in range(6)] # x1, x2, ..., x6
def noisy_hartmann_evaluation_function(parameterization):
x = np.array([parameterization.get(p_name) for p_name in param_names])
noise1, noise2 = np.random.normal(0, noise_sd, 2)
return {
"hartmann6": (hartmann6(x) + noise1, noise_sd),
"l2norm": (np.sqrt((x ** 2).sum()) + noise2, noise_sd)
}
ax_client = AxClient()
ax_client.create_experiment(
name="test_visualizations",
parameters=[
{
"name": p_name,
"type": "range",
"bounds": [0.0, 1.0],
}
for p_name in param_names
],
objective_name="hartmann6",
minimize=True,
outcome_constraints=["l2norm <= 1.25"]
)
[INFO 02-02 05:17:40] ax.service.ax_client: Starting optimization with verbose logging. To disable logging, set the `verbose_logging` argument to `False`. Note that float values in the logs are rounded to 6 decimal points.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x1. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x2. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x3. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x4. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x5. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Inferred value type of ParameterType.FLOAT for parameter x6. If that is not the expected value type, you can explicity specify 'value_type' ('int', 'float', 'bool' or 'str') in parameter dict.
[INFO 02-02 05:17:40] ax.service.utils.instantiation: Created search space: SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x3', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x4', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x5', parameter_type=FLOAT, range=[0.0, 1.0]), RangeParameter(name='x6', parameter_type=FLOAT, range=[0.0, 1.0])], parameter_constraints=[]).
[INFO 02-02 05:17:40] ax.modelbridge.dispatch_utils: Using Bayesian optimization since there are more ordered parameters than there are categories for the unordered categorical parameters.
[INFO 02-02 05:17:40] ax.modelbridge.dispatch_utils: Using Bayesian Optimization generation strategy: GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 12 trials, GPEI for subsequent trials]). Iterations after 12 will take longer to generate due to model-fitting.
for i in range(20):
parameters, trial_index = ax_client.get_next_trial()
# Local evaluation here can be replaced with deployment to external system.
ax_client.complete_trial(trial_index=trial_index, raw_data=noisy_hartmann_evaluation_function(parameters))
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 0 with parameters {'x1': 0.600227, 'x2': 0.980317, 'x3': 0.942414, 'x4': 0.476665, 'x5': 0.365389, 'x6': 0.274512}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 0 with data: {'hartmann6': (-0.651945, 0.1), 'l2norm': (1.54123, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 1 with parameters {'x1': 0.986634, 'x2': 0.117221, 'x3': 0.255546, 'x4': 0.949158, 'x5': 0.428006, 'x6': 0.011992}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 1 with data: {'hartmann6': (0.048462, 0.1), 'l2norm': (1.601498, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 2 with parameters {'x1': 0.576384, 'x2': 0.850615, 'x3': 0.584337, 'x4': 0.069568, 'x5': 0.721355, 'x6': 0.902079}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 2 with data: {'hartmann6': (0.009956, 0.1), 'l2norm': (1.733888, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 3 with parameters {'x1': 0.15922, 'x2': 0.645179, 'x3': 0.222132, 'x4': 0.635251, 'x5': 0.999075, 'x6': 0.104096}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 3 with data: {'hartmann6': (-0.653819, 0.1), 'l2norm': (1.699097, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 4 with parameters {'x1': 0.104099, 'x2': 0.933944, 'x3': 0.42663, 'x4': 0.916073, 'x5': 0.517064, 'x6': 0.573698}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 4 with data: {'hartmann6': (-0.097823, 0.1), 'l2norm': (1.704236, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 5 with parameters {'x1': 0.103375, 'x2': 0.327549, 'x3': 0.009723, 'x4': 0.913104, 'x5': 0.45098, 'x6': 0.766505}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 5 with data: {'hartmann6': (-0.134327, 0.1), 'l2norm': (1.286663, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 6 with parameters {'x1': 0.6735, 'x2': 0.607296, 'x3': 0.471511, 'x4': 0.456915, 'x5': 0.524769, 'x6': 0.130378}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 6 with data: {'hartmann6': (-0.534714, 0.1), 'l2norm': (1.226419, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 7 with parameters {'x1': 0.057797, 'x2': 0.04646, 'x3': 0.770787, 'x4': 0.226627, 'x5': 0.42694, 'x6': 0.4162}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 7 with data: {'hartmann6': (-1.006067, 0.1), 'l2norm': (0.878384, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 8 with parameters {'x1': 0.793814, 'x2': 0.796011, 'x3': 0.839933, 'x4': 0.663749, 'x5': 0.833131, 'x6': 0.999615}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 8 with data: {'hartmann6': (-0.078847, 0.1), 'l2norm': (1.9477, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 9 with parameters {'x1': 0.254719, 'x2': 0.942122, 'x3': 0.951827, 'x4': 0.069089, 'x5': 0.48823, 'x6': 0.775171}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 9 with data: {'hartmann6': (-0.19214, 0.1), 'l2norm': (1.456251, 0.1)}.
[INFO 02-02 05:17:40] ax.service.ax_client: Generated new trial 10 with parameters {'x1': 0.57471, 'x2': 0.010099, 'x3': 0.264668, 'x4': 0.993241, 'x5': 0.756086, 'x6': 0.57508}.
[INFO 02-02 05:17:40] ax.service.ax_client: Completed trial 10 with data: {'hartmann6': (0.116854, 0.1), 'l2norm': (1.334987, 0.1)}.
[INFO 02-02 05:17:41] ax.service.ax_client: Generated new trial 11 with parameters {'x1': 0.516953, 'x2': 0.757798, 'x3': 0.869154, 'x4': 0.428929, 'x5': 0.854901, 'x6': 0.992538}.
[INFO 02-02 05:17:41] ax.service.ax_client: Completed trial 11 with data: {'hartmann6': (-0.085215, 0.1), 'l2norm': (1.858812, 0.1)}.
[WARNING 02-02 05:17:41] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:09] ax.service.ax_client: Generated new trial 12 with parameters {'x1': 0.094718, 'x2': 0.123586, 'x3': 0.771925, 'x4': 0.222901, 'x5': 0.436416, 'x6': 0.304467}.
[INFO 02-02 05:18:09] ax.service.ax_client: Completed trial 12 with data: {'hartmann6': (-0.755042, 0.1), 'l2norm': (0.927186, 0.1)}.
[WARNING 02-02 05:18:09] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:10] ax.service.ax_client: Generated new trial 13 with parameters {'x1': 0.0, 'x2': 0.0, 'x3': 0.826711, 'x4': 0.168132, 'x5': 0.443425, 'x6': 0.471613}.
[INFO 02-02 05:18:10] ax.service.ax_client: Completed trial 13 with data: {'hartmann6': (-0.793586, 0.1), 'l2norm': (0.923366, 0.1)}.
[WARNING 02-02 05:18:10] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:12] ax.service.ax_client: Generated new trial 14 with parameters {'x1': 0.097057, 'x2': 0.030648, 'x3': 0.618433, 'x4': 0.306787, 'x5': 0.3996, 'x6': 0.480984}.
[INFO 02-02 05:18:12] ax.service.ax_client: Completed trial 14 with data: {'hartmann6': (-2.122568, 0.1), 'l2norm': (0.777827, 0.1)}.
[WARNING 02-02 05:18:12] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:20] ax.service.ax_client: Generated new trial 15 with parameters {'x1': 0.086156, 'x2': 0.0, 'x3': 0.573069, 'x4': 0.319927, 'x5': 0.379508, 'x6': 0.530014}.
[INFO 02-02 05:18:20] ax.service.ax_client: Completed trial 15 with data: {'hartmann6': (-2.324405, 0.1), 'l2norm': (1.009213, 0.1)}.
[WARNING 02-02 05:18:20] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:24] ax.service.ax_client: Generated new trial 16 with parameters {'x1': 0.105353, 'x2': 0.0, 'x3': 0.514095, 'x4': 0.314386, 'x5': 0.427638, 'x6': 0.506605}.
[INFO 02-02 05:18:25] ax.service.ax_client: Completed trial 16 with data: {'hartmann6': (-2.133562, 0.1), 'l2norm': (0.98663, 0.1)}.
[WARNING 02-02 05:18:25] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:37] ax.service.ax_client: Generated new trial 17 with parameters {'x1': 0.115907, 'x2': 0.0, 'x3': 0.563254, 'x4': 0.390378, 'x5': 0.32247, 'x6': 0.537632}.
[INFO 02-02 05:18:37] ax.service.ax_client: Completed trial 17 with data: {'hartmann6': (-2.257567, 0.1), 'l2norm': (0.900253, 0.1)}.
[WARNING 02-02 05:18:37] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:55] ax.service.ax_client: Generated new trial 18 with parameters {'x1': 0.069867, 'x2': 0.024548, 'x3': 0.550975, 'x4': 0.276616, 'x5': 0.289434, 'x6': 0.56398}.
[INFO 02-02 05:18:55] ax.service.ax_client: Completed trial 18 with data: {'hartmann6': (-2.942139, 0.1), 'l2norm': (0.878138, 0.1)}.
[WARNING 02-02 05:18:55] ax.utils.common.kwargs: `<class 'ax.modelbridge.torch.TorchModelBridge'>` expected argument `transform_configs` to be of type typing.Union[typing.Dict[str, typing.Dict[str, typing.Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, typing.Dict[str, typing.Any], NoneType]]], NoneType]. Got {'Winsorize': {'optimization_config': OptimizationConfig(objective=Objective(metric_name="hartmann6", minimize=True), outcome_constraints=[OutcomeConstraint(l2norm <= 1.25)])}} (type: <class 'dict'>).
[INFO 02-02 05:18:57] ax.service.ax_client: Generated new trial 19 with parameters {'x1': 0.05662, 'x2': 0.045792, 'x3': 0.52059, 'x4': 0.231451, 'x5': 0.238477, 'x6': 0.589865}.
[INFO 02-02 05:18:57] ax.service.ax_client: Completed trial 19 with data: {'hartmann6': (-2.677228, 0.1), 'l2norm': (0.796882, 0.1)}.
The plot below shows the response surface for hartmann6 metric as a function of the x1, x2 parameters.
The other parameters are fixed in the middle of their respective ranges, which in this example is 0.5 for all of them.
# this could alternately be done with `ax.plot.contour.plot_contour`
render(ax_client.get_contour_plot(param_x="x1", param_y="x2", metric_name='hartmann6'))
[INFO 02-02 05:18:57] ax.service.ax_client: Retrieving contour plot with parameter 'x1' on X-axis and 'x2' on Y-axis, for metric 'hartmann6'. Remaining parameters are affixed to the middle of their range.
The plot below allows toggling between different pairs of parameters to view the contours.
model = ax_client.generation_strategy.model
render(interact_contour(model=model, metric_name='hartmann6'))
This plot illustrates the tradeoffs achievable for 2 different metrics. The plot takes the x-axis metric as input (usually the objective) and allows toggling among all other metrics for the y-axis.
This is useful to get a sense of the pareto frontier (i.e. what is the best objective value achievable for different bounds on the constraint)
render(plot_objective_vs_constraints(model, 'hartmann6', rel=False))
CV plots are useful to check how well the model predictions calibrate against the actual measurements. If all points are close to the dashed line, then the model is a good predictor of the real data.
cv_results = cross_validate(model)
render(interact_cross_validation(cv_results))
Slice plots show the metric outcome as a function of one parameter while fixing the others. They serve a similar function as contour plots.
render(plot_slice(model, "x2", "hartmann6"))
Tile plots are useful for viewing the effect of each arm.
render(interact_fitted(model, rel=False))
Total runtime of script: 1 minutes, 40.64 seconds.